from typing import Dict, Any, AsyncGenerator
import logging
import httpx
import asyncio
from .apis.ds_ggb_generation_api import GGBGenerationAPI
from .apis.ds_html_generation_api import HTMLGenerationAPI
from .apis.ds_element_extraction_api import TwoDElementExtractionAPI

# 导入函数API
from .function_api.function_element_extraction_api import FunctionElementExtractionAPI
from .function_api.function_html_generation_api import FunctionHTMLGenerationAPI

# 导入知识点API
from .konwledge_api.knowledge_element_extraction_api import KnowledgeElementExtractionAPI
from .konwledge_api.knowledge_html_generation_api import KnowledgeHTMLGenerationAPI

# 导入3D API（3D的元素提取使用现有的API接口调用）
import sys
import os
# 手动导入3D APIs
threed_apis_path = os.path.join(os.path.dirname(__file__), "3D_apis")
sys.path.append(threed_apis_path)

from threed_ggb_generation_api import ThreeDGGBGenerationAPI
from threed_html_generation_api import ThreeDHTMLGenerationAPI

class GGBManagement:
    def __init__(self):
        # 2D几何APIs
        self.twod_element_extraction_api = TwoDElementExtractionAPI()
        self.ggb_generation_api = GGBGenerationAPI()
        self.html_generation_api = HTMLGenerationAPI()
        
        # 函数APIs
        self.function_element_extraction_api = FunctionElementExtractionAPI()
        self.function_html_generation_api = FunctionHTMLGenerationAPI()
        
        # 知识点APIs
        self.knowledge_element_extraction_api = KnowledgeElementExtractionAPI()
        self.knowledge_html_generation_api = KnowledgeHTMLGenerationAPI()
        
        # 3D几何APIs（元素提取通过HTTP调用现有服务）
        self.threed_ggb_generation_api = ThreeDGGBGenerationAPI()
        self.threed_html_generation_api = ThreeDHTMLGenerationAPI()
        
        logging.basicConfig(
            filename="ggb_management.log",
            level=logging.INFO,
            format="%(asctime)s - %(levelname)s - %(message)s"
        )

    async def process_problem(self, problem: str) -> AsyncGenerator[Dict[str, Any], None]:
        """
        Process a math problem through the complete pipeline:
        1. Extract elements
        2. Generate GGB commands
        3. Generate HTML
        
        Args:
            problem: The math problem text
            
        Yields:
            Dictionary containing the type and content of each step's output
        """
        try:
            logging.info(f"开始处理问题: {problem[:100]}...")
            
            # Step 1: Extract elements using local API (streaming)
            logging.info("开始元素提取（使用本地API）...")
            
            elements_content = ""
            async for chunk in self.twod_element_extraction_api.extract_elements_stream(problem):
                if chunk["type"] == "reasoning":
                    yield {
                        "type": "reasoning",
                        "content": chunk["content"],
                        "step": "element_extraction"
                    }
                elif chunk["type"] == "element":
                    elements_content += chunk["content"]
                    yield {
                        "type": "element",
                        "content": chunk["content"],
                        "step": "element_extraction"
                    }
                elif chunk["type"] == "error":
                    yield {
                        "type": "error",
                        "content": chunk["content"],
                        "step": "element_extraction_error"
                    }
                    return
            
            elements = elements_content
            if elements and len(elements.strip()) > 3:  # 确保不是只有"###"
                logging.info(f"2D元素提取成功，内容: {elements[:200]}...")
                
                # Step 2: Generate GGB commands
                logging.info("开始生成GGB命令...")
                ggb_result = await self.ggb_generation_api.generate_ggb(elements)
                logging.info(f"GGB生成结果: {ggb_result}")
                
                yield {
                    "type": "ggb_commands",
                    "content": ggb_result["commands"],
                    "step": "ggb_generation"
                }
                
                # Step 3: Generate HTML
                logging.info("开始生成HTML...")
                html_result = await self.html_generation_api.generate_html(ggb_result["commands"])
                logging.info(f"HTML生成完成，长度: {len(html_result['html_code'])}")
                
                yield {
                    "type": "html_code",
                    "content": html_result["html_code"],
                    "step": "html_generation"
                }
            else:
                logging.error(f"2D元素提取返回的内容无效: '{elements}'")
                yield {
                    "type": "error",
                    "content": f"元素提取失败，返回内容无效: '{elements}'",
                    "step": "element_extraction_error"
                }
                return
            
            # Signal completion
            logging.info("所有步骤处理完成")
            yield {
                "type": "complete",
                "content": "Processing completed",
                "step": "complete"
            }
            
        except Exception as e:
            logging.error(f"Error in process_problem: {str(e)}")
            yield {
                "type": "error",
                "content": str(e),
                "step": "error"
            }

    async def process_function_problem(self, problem: str) -> AsyncGenerator[Dict[str, Any], None]:
        """
        Process a function problem through the complete pipeline:
        1. Extract function elements (streaming)
        2. Generate HTML
        
        Args:
            problem: The function problem text
            
        Yields:
            Dictionary containing the type and content of each step's output
        """
        try:
            logging.info(f"开始处理函数问题: {problem[:100]}...")
            
            # Step 1: Extract function elements (streaming)
            logging.info("开始函数元素提取（流式）...")
            
            elements_content = ""
            async for chunk in self.function_element_extraction_api.extract_elements_stream(problem):
                if chunk["type"] == "reasoning":
                    yield {
                        "type": "reasoning",
                        "content": chunk["content"],
                        "step": "function_element_extraction"
                    }
                elif chunk["type"] == "element":
                    elements_content += chunk["content"]
                    yield {
                        "type": "element",
                        "content": chunk["content"],
                        "step": "function_element_extraction"
                    }
                elif chunk["type"] == "error":
                    yield {
                        "type": "error",
                        "content": chunk["content"],
                        "step": "function_element_extraction_error"
                    }
                    return
            
            if elements_content and len(elements_content.strip()) > 3:
                logging.info(f"函数元素提取成功，内容: {elements_content[:200]}...")
                
                # Step 2: Generate HTML
                logging.info("开始生成函数HTML...")
                html_result = await self.function_html_generation_api.generate_html(elements_content)
                logging.info(f"函数HTML生成完成，长度: {len(html_result['html_code'])}")
                
                yield {
                    "type": "html_code",
                    "content": html_result["html_code"],
                    "step": "function_html_generation"
                }
            else:
                logging.error(f"函数元素提取返回的内容无效: '{elements_content}'")
                yield {
                    "type": "error",
                    "content": f"函数元素提取失败，返回内容无效: '{elements_content}'",
                    "step": "function_element_extraction_error"
                }
                return
            
            # Signal completion
            logging.info("函数问题处理完成")
            yield {
                "type": "complete",
                "content": "Function processing completed",
                "step": "complete"
            }
            
        except Exception as e:
            logging.error(f"Error in process_function_problem: {str(e)}")
            yield {
                "type": "error",
                "content": str(e),
                "step": "error"
            }

    async def process_knowledge_problem(self, knowledge_point: str) -> AsyncGenerator[Dict[str, Any], None]:
        """
        Process a knowledge point through the complete pipeline:
        1. Extract knowledge elements (streaming)
        2. Generate HTML
        
        Args:
            knowledge_point: The knowledge point text
            
        Yields:
            Dictionary containing the type and content of each step's output
        """
        try:
            logging.info(f"开始处理知识点问题: {knowledge_point[:100]}...")
            
            # Step 1: Extract knowledge elements (streaming)
            logging.info("开始知识点元素提取（流式）...")
            
            elements_content = ""
            async for chunk in self.knowledge_element_extraction_api.extract_elements_stream(knowledge_point):
                if chunk["type"] == "reasoning":
                    yield {
                        "type": "reasoning",
                        "content": chunk["content"],
                        "step": "knowledge_element_extraction"
                    }
                elif chunk["type"] == "element":
                    elements_content += chunk["content"]
                    yield {
                        "type": "element",
                        "content": chunk["content"],
                        "step": "knowledge_element_extraction"
                    }
                elif chunk["type"] == "error":
                    yield {
                        "type": "error",
                        "content": chunk["content"],
                        "step": "knowledge_element_extraction_error"
                    }
                    return
            
            if elements_content and len(elements_content.strip()) > 3:
                logging.info(f"知识点元素提取成功，内容: {elements_content[:200]}...")
                
                # Step 2: Generate HTML
                logging.info("开始生成知识点HTML...")
                html_result = await self.knowledge_html_generation_api.generate_html(elements_content)
                logging.info(f"知识点HTML生成完成，长度: {len(html_result['html_code'])}")
                
                yield {
                    "type": "html_code",
                    "content": html_result["html_code"],
                    "step": "knowledge_html_generation"
                }
            else:
                logging.error(f"知识点元素提取返回的内容无效: '{elements_content}'")
                yield {
                    "type": "error",
                    "content": f"知识点元素提取失败，返回内容无效: '{elements_content}'",
                    "step": "knowledge_element_extraction_error"
                }
                return
            
            # Signal completion
            logging.info("知识点问题处理完成")
            yield {
                "type": "complete",
                "content": "Knowledge processing completed",
                "step": "complete"
            }
            
        except Exception as e:
            logging.error(f"Error in process_knowledge_problem: {str(e)}")
            yield {
                "type": "error",
                "content": str(e),
                "step": "error"
            }

    async def process_threed_problem(self, problem: str, image_path: str = None) -> AsyncGenerator[Dict[str, Any], None]:
        """
        Process a 3D geometry problem through the complete pipeline:
        1. Extract 3D elements (streaming, with multimodal support)
        2. Generate GGB commands
        3. Generate HTML
        
        Args:
            problem: The 3D geometry problem text
            image_path: Optional path to image file
            
        Yields:
            Dictionary containing the type and content of each step's output
        """
        try:
            logging.info(f"开始处理3D几何问题: {problem[:100]}...")
            
            # Step 1: Extract 3D elements (streaming)
            logging.info("开始3D元素提取（流式，多模态）...")
            
            elements_content = ""
            # Note: 3D API的元素提取使用多模态模型，这里需要调用现有的API
            try:
                if image_path:
                    # 使用图像提取接口
                    result = await self.call_threed_element_extraction_api_with_image(problem, image_path)
                else:
                    # 使用文字提取接口
                    result = await self.call_threed_element_extraction_api(problem)
                
                # 输出reasoning内容（分块输出以模拟流式）
                reasoning = result.get("reasoning", "")
                if reasoning:
                    chunk_size = 100
                    for i in range(0, len(reasoning), chunk_size):
                        chunk = reasoning[i:i+chunk_size]
                        yield {
                            "type": "reasoning",
                            "content": chunk,
                            "step": "threed_element_extraction"
                        }
                        await asyncio.sleep(0.01)
                
                elements_content = result.get("elements", "")
                if elements_content and len(elements_content.strip()) > 3:
                    logging.info(f"3D元素提取成功，内容: {elements_content[:200]}...")
                    
                    yield {
                        "type": "element",
                        "content": elements_content,
                        "step": "threed_element_extraction"
                    }
                    
                    # Step 2: Generate GGB commands
                    logging.info("开始生成3D GGB命令...")
                    ggb_result = await self.threed_ggb_generation_api.generate_ggb(elements_content)
                    logging.info(f"3D GGB生成结果: {ggb_result}")
                    
                    yield {
                        "type": "ggb_commands",
                        "content": ggb_result["commands"],
                        "step": "threed_ggb_generation"
                    }
                    
                    # Step 3: Generate HTML
                    logging.info("开始生成3D HTML...")
                    html_result = await self.threed_html_generation_api.generate_html(ggb_result["commands"])
                    logging.info(f"3D HTML生成完成，长度: {len(html_result['html_code'])}")
                    
                    yield {
                        "type": "html_code",
                        "content": html_result["html_code"],
                        "step": "threed_html_generation"
                    }
                else:
                    logging.error(f"3D元素提取返回的内容无效: '{elements_content}'")
                    yield {
                        "type": "error",
                        "content": f"3D元素提取失败，返回内容无效: '{elements_content}'",
                        "step": "threed_element_extraction_error"
                    }
                    return
                    
            except Exception as api_error:
                logging.error(f"调用3D API失败: {api_error}")
                yield {
                    "type": "error",
                    "content": f"3D元素提取API调用失败: {str(api_error)}",
                    "step": "threed_api_call_error"
                }
                return
            
            # Signal completion
            logging.info("3D几何问题处理完成")
            yield {
                "type": "complete",
                "content": "3D processing completed",
                "step": "complete"
            }
            
        except Exception as e:
            logging.error(f"Error in process_threed_problem: {str(e)}")
            yield {
                "type": "error",
                "content": str(e),
                "step": "error"
            }

    async def call_threed_element_extraction_api(self, problem: str) -> Dict[str, str]:
        """
        调用3D元素提取API（非流式）
        """
        try:
            async with httpx.AsyncClient(timeout=300) as client:
                response = await client.post(
                    "http://localhost:8001/extract_elements",
                    json={"problem": problem, "problem_img": ""}
                )
                if response.status_code == 200:
                    result = response.json()
                    logging.info(f"3D API调用成功，reasoning长度: {len(result.get('reasoning', ''))}, elements长度: {len(result.get('elements', ''))}")
                    return result
                else:
                    logging.error(f"3D API调用失败，状态码: {response.status_code}, 响应: {response.text}")
                    raise Exception(f"3D API调用失败: {response.status_code}")
        except Exception as e:
            logging.error(f"调用3D API异常: {e}")
            raise

    async def call_threed_element_extraction_api_with_image(self, problem: str, image_path: str) -> Dict[str, str]:
        """
        调用3D元素提取API（包含图像，非流式）
        """
        try:
            async with httpx.AsyncClient(timeout=300) as client:
                response = await client.post(
                    "http://localhost:8001/extract_elements_image",
                    json={"problem": problem, "problem_img": image_path}
                )
                if response.status_code == 200:
                    result = response.json()
                    logging.info(f"3D图像API调用成功，reasoning长度: {len(result.get('reasoning', ''))}, elements长度: {len(result.get('elements', ''))}")
                    return result
                else:
                    logging.error(f"3D图像API调用失败，状态码: {response.status_code}, 响应: {response.text}")
                    raise Exception(f"3D图像API调用失败: {response.status_code}")
        except Exception as e:
            logging.error(f"调用3D图像API异常: {e}")
            raise
